import os
import torch
from torch.utils.data import DataLoader

from tqdm import tqdm 
import matplotlib.pyplot as plt

from model import get_crnn,get_res
from dataset import SunDataset
from utils import AverageMeter

from visdom import Visdom
import datetime as dt


device='cuda:0'
def main():
    if not os.path.exists('logs/'):
        os.makedirs('logs/')
    vis = Visdom(env = 'sun_val',log_to_filename='logs/sun_val_{}.log'.format(
                    dt.datetime.now().strftime('%Y-%m-%d_T%H-%M-%S')))
    assert vis.check_connection()
    checkpoint = torch.load('saved_models/res18_best.pth')['model']
    # checkpoint = torch.load('saved_models/model_last.pth')

    model = get_res()
    model.load_state_dict(checkpoint)
    model.to(device)

    criterion = torch.nn.L1Loss()
    loss_meter = AverageMeter()

    val_set = SunDataset(root='./data', split='val', size=224, gray=False)
    val_loader = DataLoader(val_set,batch_size=1,shuffle=False,num_workers=0)

    gts = []
    preds = []

    model.eval()
    bar = tqdm(total=len(val_loader))
    with torch.no_grad():
        for i, (x,y) in enumerate(val_loader):

            y_hat = model(x.to(device)).cpu()
            gts.append(y.mean().item())
            preds.append(y_hat.mean().item())

            batch_size = x.shape[0]
            loss = criterion(y_hat, y)
            loss_meter.update(loss.item()/batch_size)

            vis.line([[y, y_hat]], [i], win='pred', opts={'legend':['gt','pred']}, update='append')
            vis.line([loss_meter.avg], [i], win='loss', update='append')
            bar.update()
            bar.set_description('loss:{:.5f}'.format(loss_meter.avg))

            # if i == 500:
            #     break
        bar.close()
    plt.figure()
    plt.plot(gts)
    plt.plot(preds)
    plt.legend(['gts','preds'])
    plt.figure()
    plt.plot(preds)
    plt.show()
    
 

if __name__ == "__main__":
    main()